﻿using System.Collections.Generic;

using Vestras.StarCraft2.Luna.Parser.Ast;

namespace Vestras.StarCraft2.Luna.Validation {
    internal static class ScopeEntityFilterer {
        public static IEnumerable<Entity> FilterEntities(IEnumerable<Entity> entities, Entity parent) {
            if (entities == null) {
                return new List<Entity>();
            }

            IList<Entity> list = new List<Entity>(entities);
            if (list.Count == 1) {
                return list;
            }

            Function function = parent is Function ? parent as Function : parent.GetLogicalParentOfType<Function>();
            if (function == null) {
                return list;
            }

            IList<Entity> containedEntities = new List<Entity>();
            IDictionary<Entity, Entity> parentMaps = new Dictionary<Entity, Entity>();
            foreach (Entity entity in list) {
                if (function.Children.Contains(entity)) {
                    containedEntities.Add(entity);
                }

                parentMaps.Add(entity, entity.Parent);
            }

            if (containedEntities.Count == 0) {
                List<Entity> entitiesWithNoParent = new List<Entity>();
                foreach (KeyValuePair<Entity, Entity> pair in parentMaps) {
                    if (pair.Value == null) {
                        entitiesWithNoParent.Add(pair.Key);
                    }
                }

                if (entitiesWithNoParent.Count > 0) {
                    entitiesWithNoParent.ForEach(item => containedEntities.Add(item));
                } else {
                    foreach (KeyValuePair<Entity, Entity> pair in parentMaps) {
                        containedEntities.Add(pair.Key);
                    }
                }
            }

            if (containedEntities.Count == 1) {
                list = containedEntities;
            }

            return list;
        }
    }
}
